"""
=============================================================
        LangChain Few-Shot Prompt 示例（语义相似选择器）
=============================================================

功能：
    - 根据输入问题，从已有问答示例中选择最语义相似的示例。
    - 使用 HuggingFace 嵌入模型 + Chroma 向量数据库进行语义检索。
    - 构建 FewShotPromptTemplate，将选出的示例拼接成完整 prompt。
用途：
    - 少样本问答（Few-Shot Q&A）
    - 动态生成 prompt，提高大模型回答质量
"""

# -------------------------------
# 导入依赖
# -------------------------------
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain_community.vectorstores import Chroma  # 本地向量数据库
from langchain_huggingface import HuggingFaceEmbeddings  # 文本向量嵌入模型
from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate  # Prompt 模板


# -------------------------------
# 示例数据集：定义几个问答示例
# -------------------------------
examples = [
    {
        "question": "谁的寿命更长，穆罕默德·阿里还是艾伦·图灵",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：穆罕默德·阿里去世时多大？
            中间答案：穆罕默德·阿里去世时74岁。
            根据：艾伦·图灵去世时41岁。
            所以最终答案是：穆罕默德·阿里
        """
    },
    {
        "question": "肯德基的创始人是什么时候出生的？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：肯德基的创始人是谁？
            中间答案：肯德基由肯德基上校创立。
            跟进：肯德基上校是什么时候出生的？
            中间答案：肯德基上校于1933年4月2日出生。
            所以最终答案是：1933年4月2日。
        """
    },
    {
        "question": "《大白鲨》和《皇家赌场》的导演都来自同一个国家吗？",
        "answer": """
            这里需要跟进问题吗：是的。
            跟进：《大白鲨》的导演是谁？
            中间答案：《大白鲨》的导演是Steven Spielberg。
            跟进：Steven Spielberg来自哪里？
            中间答案：美国。
            跟进：《皇家赌场》的导演是谁？
            中间答案：《皇家赌场》的导演是Martin Campbell。
            跟进：Martin Campbell来自哪里？
            中间答案：新西兰。
            所以最终答案是：不是
        """
    }
]


# -------------------------------
# 定义每个示例的模板
# -------------------------------
example_template = PromptTemplate(
    input_variables=["question", "answer"],
    template="问题：{question}\n答案：{answer}"
)


# -------------------------------
# 创建语义相似选择器
# -------------------------------
# 1️⃣ 创建 HuggingFace 嵌入模型
#embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
embeddings = HuggingFaceEmbeddings(model_name="D:/models/BAAIbge-base-zh-v1.5", model_kwargs={'device': 'cpu'})


# 2️⃣ 使用 SemanticSimilarityExampleSelector 自动选择最相似示例
example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,   # 示例列表
    embeddings, # 用于生成向量
    Chroma,     # 向量数据库
    k=1         # 每次选择 1 个最相似示例
)


# -------------------------------
# 构建 Few-Shot Prompt
# -------------------------------
few_shot_prompt = FewShotPromptTemplate(
    example_selector=example_selector,   # 选择器
    example_prompt=example_template,     # 每个示例的格式
    input_variables=["question"],        # 新问题变量
    prefix="下面是一些问答示例：",       # prompt 前缀
    suffix="请根据示例回答下一个问题：{question}"  # prompt 后缀
)


# -------------------------------
# 使用示例：生成最终 prompt
# -------------------------------
question = "乔治·华盛顿的父亲是谁？"
final_prompt = few_shot_prompt.format(question=question)

# 打印生成的 prompt
print("生成的 Few-Shot Prompt：\n")
print(final_prompt)
